"""
Author: xubing
Date: 2024-01-12 06:31:46
LastEditors: xubing
LastEditTime: 2024-03-07 09:32:07
Description: 
"""

from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableMap


class RAG:
    def __init__(self, retriever, llm) -> None:
        self.template = """
        你是一位证券基金的专家,你的任务是帮助用户解答关于证券基金的各种问题，包括市场动态、基金选择、投资策略等。
        你的回答基于最新的金融数据和智能算法，旨在为用户提供准确、及时的投资建议。
        你的任务是帮助用户快速理解各类理财产品的特性、风险与回报。
        - 你能迅速理解用户想要问的问题
        - 你能迅速从可参考信息中提炼出关键信息
        - 你能对基金的历史表现、风险等级等进行深度分析。
        - 如果无法通过上下文得出问题的答案, 请如实回答.
        - 可以对答案进行润色, 但请不要编造答案.
        
        
        用户问题: {question}
        以下是从理财基金报告中提取出的参考信息: {context}
        
        根据参考信息来回复用户的问题.
        """
        self.retriever = retriever
        self.llm = llm

    def get_retrival_result(self, user_message):
        # 创建prompt模板

        # 由模板生成prompt
        prompt = ChatPromptTemplate.from_template(self.template)

        # 创建chain
        chain = (
            RunnableMap(
                {
                    # "example": lambda x: x["example"],
                    "question": lambda x: x["question"],
                    "context": lambda x: self.retriever.get_relevant_documents(
                        x["question"]
                    ),
                }
            )
            | prompt
            | self.llm
            | StrOutputParser()
        )
        relevant_doc = self.retriever.get_relevant_documents(user_message)
        response = chain.invoke({"question": user_message})
        return relevant_doc, response


if __name__ == "__main__":
    import yaml

    from fin_embedding import MilvusStore

    with open("config.yaml", "r") as f:
        config = yaml.safe_load(f)
    Config = None

    milvus_store = MilvusStore(
        embedding_model_name_or_path=config.get("embedding_model_name_or_path"),
        milvus_conn_args=config.get("milvus_conn_args"),
    )
    collection_name = "fin_emb"
    retriever = milvus_store.get_retriever(collection_name)

    from langchain_openai import AzureChatOpenAI, ChatOpenAI

    lc_az_gpt35 = AzureChatOpenAI(**config.get("gpt35"))

    # prompt = ""
    print(
        RAG(retriever, lc_az_gpt35).get_retrival_result("基金的合同生效日是什么时候?")
    )
